from random import shuffle

import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

from config import parser

args = parser.parse_args()

'''
1. 对图片进行按比例缩放
2. 对图片进行随机位置的截取
3. 对图片进行随机的水平和竖直翻转
4. 对图片进行随机角度的旋转
5. 对图片进行亮度、对比度和颜色的随机变化


数据增强最新的方式
1、Mixup
2、Cutout
3、Cutmix
4、Mosaic

'''
data_transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.RandomRotation(30),
    transforms.ColorJitter(brightness=0.5, contrast=0.5, hue=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
data_transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


# 自己写Dataset至少需要有这样的格式
class Dataset(Dataset):
    def __init__(self, lines, type="train"):
        super(Dataset, self).__init__()
        self.base_path = args.base_data_path
        self.annotation_lines = lines
        self.train_batches = len(self.annotation_lines)
        self.type = type

    def __len__(self):
        return self.train_batches

    def __getitem__(self, index):
        if index == 0:
            shuffle(self.annotation_lines)
        n = len(self.annotation_lines)
        index = index % n
        img, y = self.collect_image_label(self.annotation_lines[index])
        h, w = img.height, img.width
        if h / w > 1:
            img.resize((int(224 * (w / h)), 224), Image.BICUBIC)
        else:
            img.resize((224, int(224 * (h / w))), Image.BICUBIC)

        new_img = Image.new("RGB", (224, 224), 0)
        new_img.paste(img, (0, 0))
        if self.type == "train":
            new_img = data_transform_train(new_img)
        else:
            new_img = data_transform_test(new_img)

        temp_y = int(y)
        return new_img, temp_y

    def collect_image_label(self, line):

        line = line.split(' ')
        image_path = line[1].strip()
        label = line[0]
        image = Image.open(image_path.replace("G:/datasets/分类/dog_class_num=10","/media/luotianhang/My Passport/datasets/分类/dog_class_num=10").replace("\\","/")).convert("RGB")
        # image = Image.open(image).convert("RGB")
        return image, label

    def collect_image_label_cat_dog(self, line):
        image = Image.open(
            line.strip().replace("E:/Datasets2/cat_dog/sort/ImageNet", "/workspace/dog10_classify/Images")).convert(
            "RGB")
        if "dog/dog" in line:
            return image, 0
        else:
            return image, 1

    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    def img_augment(self, image):

        # # 随机位置裁剪
        # random_crop = self.rand() < 0.5
        # # 中心裁剪
        # center_crop = self.rand() < 0.5
        # # 填充后随机裁剪
        # random_crop_padding = self.rand() < 0.5
        # 水平翻转
        h_flip = self.rand() < 0.5
        # 竖直翻转
        v_flip = self.rand() < 0.5
        # # 亮度
        # bright = self.rand() < 0.5
        # # 对比度
        # contrast = self.rand() < 0.5
        # # 饱和度
        # saturation = self.rand() < 0.5
        # # 颜色随机变换
        # color = self.rand() < 0.5
        # compose = self.rand() < 0.5
        # # 旋转30
        # rotate = self.rand() < 0.5

        if h_flip:
            image = transforms.RandomHorizontalFlip()(image)
        if v_flip:
            image = transforms.RandomVerticalFlip()(image)
        # if rotate:
        #     image = transforms.RandomRotation(30)(image)
        # if bright:
        #     image = transforms.ColorJitter(brightness=1)(image)
        # if contrast:
        #     image = transforms.ColorJitter(contrast=1)(image)
        # if saturation:
        #     image = transforms.ColorJitter(saturation=1)(image)
        # if color:
        #     image = transforms.ColorJitter(hue=0.5)(image)
        # if compose:
        #     image = transforms.ColorJitter(0.5, 0.5, 0.5)(image)
        # if random_crop:
        #     image = transforms.RandomCrop(100)(image)
        # if center_crop:
        #     image = transforms.CenterCrop(100)(image)
        # if random_crop_padding:
        #     image = transforms.RandomCrop(100, padding=8)(image)

        return image


'''

/workspace/dog_cat_classify/images/train/Yorkshireterrier/215.closeup-portrait-golden-yorkshire-terrier-260nw-1382906201.jpg
/workspace/dog_cat_classify/images/train/Yorkshireterrier/215.closeup-portrait-golden-yorkshire-terrier-260nw-1382906201.jpg
'''